############################
#### figure_2.R         ####
#### Generates Figure 2 ####
############################

tab_in = import(here("Data","var_data.csv"))
vars_fact = tab_in |> 
  filter(is_fact == 1) |> 
  distinct(var2) |> 
  pull(var2)

data_indiv_svy = data |> 
  mutate(pid = case_when(
    pid == "Democrat" ~ 1,
    pid == "Republican" ~ 2,
    pid == "Independent" ~ 3
  )) |> 
  mutate(across(all_of(vars_fact), as.factor)) |> 
  as_survey_design(weights = wt_adj, ids = ID_ANON)

pred_probs = map_dfr(unique(tab_in$var), \(x){
  var_in = filter(tab_in, var == x)
  # Make percentiles
  if(!is.na(unique(var_in$ptile))){
    pred_vals = data.frame(
      val_label = c("Low (25th Perc.)", "Median", "High (75th Perc.)"),
      val = quantile(data[[x]], c(.25,.5,.75), na.rm = T)
    )
  } else {
    pred_vals = distinct(var_in, val_label, val)
  }
  if(unique(var_in$is_fact == 1)){
    pred_vals = mutate(pred_vals, val = as.factor(val))
  }
  names(pred_vals)[2] = x
  # Model
  f = as.formula(paste0("affpol ~ ", x))
  mod = survey::svyglm(f, design = data_indiv_svy)
  # Predict
  predict(mod, newdata = pred_vals, type = 'response', se = T) |> 
    as_tibble() |> 
    bind_cols(pred_vals) |> 
    select(response:val_label) |> 
    left_join(var_in)
})

fig_2 = pred_probs |> 
  ggplot(aes(x = response, y = str_wrap(reorder(var_label, response), 15), group = val_label)) +
  facet_wrap(~ class, ncol = 4, scales = "free_y") +
  geom_vline(xintercept = 50, alpha = .25) +
  geom_pointrange(aes(xmin = response - 1.96*SE,
                      xmax = response + 1.96*SE),
                  position = position_dodge(0.65)) +
  geom_text(aes(label = val_label, x = response + 1.96*SE + .25), hjust = 0, size = 3,
            position = position_dodge(0.65)) +
  labs(y = NULL, x = 'Predicted Affective Polarization') +
  theme_prl() +
  theme(axis.text.y = element_text(lineheight = 0.7))

print(fig_2)

ggsave(here("Plots","figure_2.pdf"), fig_2,
       units = "in", height = 4, width = 14, dpi = 600)

#### Descriptives ####

pred_probs |> 
  group_by(class, var) |> 
  summarise(range = abs(max(response) - min(response))) |> 
  group_by(class) |> 
  summarise(mean_range = mean(range, na.rm = T))
